package com.beit.mybatis.demo.config;

import com.baomidou.mybatisplus.core.conditions.AbstractWrapper;
import com.baomidou.mybatisplus.core.conditions.SharedString;
import com.baomidou.mybatisplus.core.conditions.query.Query;
import com.baomidou.mybatisplus.core.conditions.segments.MergeSegments;
import com.baomidou.mybatisplus.core.enums.SqlKeyword;
import com.baomidou.mybatisplus.core.metadata.TableFieldInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import org.apache.commons.collections4.CollectionUtils;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class CustomQueryWrapper<T> extends AbstractWrapper<T, String, CustomQueryWrapper<T>>
        implements Query<CustomQueryWrapper<T>, T, String> {

    private final SharedString sqlSelect = new SharedString();

    public CustomQueryWrapper() {
    }

    public CustomQueryWrapper(T entity) {
        super.setEntity(entity);
        super.initNeed();
    }

    public CustomQueryWrapper(Class<T> entityClass) {
        super.setEntityClass(entityClass);
        super.initNeed();
    }

    public CustomQueryWrapper(T entity, String... columns) {
        super.setEntity(entity);
        super.initNeed();
        this.select(columns);
    }

    private CustomQueryWrapper(T entity, Class<T> entityClass, AtomicInteger paramNameSeq,
                         Map<String, Object> paramNameValuePairs, MergeSegments mergeSegments, SharedString paramAlias,
                         SharedString lastSql, SharedString sqlComment, SharedString sqlFirst) {
        super.setEntity(entity);
        super.setEntityClass(entityClass);
        this.paramNameSeq = paramNameSeq;
        this.paramNameValuePairs = paramNameValuePairs;
        this.expression = mergeSegments;
        this.paramAlias = paramAlias;
        this.lastSql = lastSql;
        this.sqlComment = sqlComment;
        this.sqlFirst = sqlFirst;
    }

    @Override
    protected CustomQueryWrapper<T> instance() {
        return new CustomQueryWrapper<>(getEntity(), getEntityClass(), paramNameSeq, paramNameValuePairs, new MergeSegments(),
                paramAlias, SharedString.emptyString(), SharedString.emptyString(), SharedString.emptyString());
    }

    @Override
    public CustomQueryWrapper<T> select(boolean condition, List<String> columns) {
        if (condition && CollectionUtils.isNotEmpty(columns)) {
            this.sqlSelect.setStringValue(String.join(StringPool.COMMA, columns));
        }
        return typedThis;
    }

    @Override
    public CustomQueryWrapper<T> select(Class<T> entityClass, Predicate<TableFieldInfo> predicate) {
        super.setEntityClass(entityClass);
        this.sqlSelect.setStringValue(TableInfoHelper.getTableInfo(getEntityClass()).chooseSelect(predicate));
        return typedThis;
    }

    public CustomQueryWrapper<T> in(List<T> inList, String... columns) {
        if (CollectionUtils.isEmpty(inList) || columns.length == 0) {
            return typedThis;
        }
        TableInfo tableInfo = TableInfoHelper.getTableInfo(getEntityClass());

        return maybeDo(true, () -> appendSqlSegments(
                () -> Stream.of(columns).map(super::columnToString)
                        .collect(Collectors.joining(StringPool.COMMA, StringPool.LEFT_BRACKET, StringPool.RIGHT_BRACKET)),
                SqlKeyword.IN,
                () -> inList.stream().map(t -> Stream.of(columns)
                                .map(col -> formatParam(null, tableInfo.getPropertyValue(t, col)))
                                .collect(Collectors.joining(StringPool.COMMA, StringPool.LEFT_BRACKET, StringPool.RIGHT_BRACKET)))
                        .collect(Collectors.joining(StringPool.COMMA, StringPool.LEFT_BRACKET, StringPool.RIGHT_BRACKET))));
    }

    public CustomQueryWrapper<T> in(Collection<Map<String, Object>> coll, String... columns) {
        return maybeDo(CollectionUtils.isNotEmpty(coll) && columns.length > 0,
                () -> appendSqlSegments(
                        () -> Stream.of(columns).map(super::columnToString)
                                .collect(Collectors.joining(StringPool.COMMA, StringPool.LEFT_BRACKET, StringPool.RIGHT_BRACKET)),
                        SqlKeyword.IN,
                        () -> coll.stream().map(map -> Stream.of(columns)
                                        .map(col -> formatParam(null, map.get(col)))
                                        .collect(Collectors.joining(StringPool.COMMA, StringPool.LEFT_BRACKET, StringPool.RIGHT_BRACKET)))
                                .collect(Collectors.joining(StringPool.COMMA, StringPool.LEFT_BRACKET, StringPool.RIGHT_BRACKET))));
    }
}
